#import ImageMetrics as IM
#import Metric
import pylab
import Image
import numpy as np

import JT
#import Calibration

class Test():
    def __init__(self):    
        """
        This function will plot a metric over the given range, for one element of 
        the pose.
        Inputs:
            eval_range: an iterable range of positions at which to calculate the metric
            eval_param: the element of the pose that will be modified over the range
        output:
            val: array of values returned by the metric at each position in eval_range
                
    
        """
    #===============================================================================
    #    Data Paths and Parameters
    #===============================================================================   
        fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
       
        inputVolFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mha'
        staFile = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
        calFileLat = 'C:/Users/bryan/bryan-code/trunk/TestData/ext_cal1.txt'
        
        debugFolder = 'C:/Users/bryan/bryan-code/trunk/Images/debug/'
        
        roi = [[210,187],[277,370]]
        body_center = (3.906, 2.17, -7.378)
        
        #Specify the scale of the matching (integer, larger = more coarse)
        level = 4
        
        # ****Make sure angle and translation ranges are the same length
        # Iterate over a range of translations from -30 to 30 mm
        numel = 10
        ang = np.pi/numel
        
        
        self.trans_range = np.arange(-numel,numel+1,1)
        # Iterate over a range of angles from -pi/4 (45deg) to pi/4 in 2.864 deg increments
        self.angle_range = np.arange(-ang,ang+ang/numel,1*ang/numel)
            
    #===============================================================================
    #    Setup DRR
    #===============================================================================  
        xraycam = JT.XRayCamera()
        drr = JT.DRR()
        drr.SetXRayCamera(xraycam)
        drr.SetBlackOnWhiteImage()
        drr.InteractionOff()
        drr._renWin.SetOffScreenRendering(1)
        
        cal = JT.ExternalCalibration()
        cal.LoadConsolidatedCalibration(calFileLat)
        cal.LoadStaFile(staFile)
        
        xraycam.SetExternalCalibration(cal)
        
        volume = JT.Volume()
        volume.SetVolumeFilename(inputVolFileName)
        volume.SetOriginalTransform(cal._VolumeT)
        volume.SetBodyCenter((3.906, 2.17, -7.378))
        drr.AddVolume(volume)
        #drr._volume.UseRayCastMapper()
        
        # Set Fixed image using path
        fixedImage = JT.FixedImage()
        fixedImage.SetFileName(fixedImageName)
        
        self.reg=JT.Registration()
        self.reg.SetFixedImage(fixedImage.GetImage(0))
        self.reg.SetMovingImageGenerator(drr)
        #self.reg._useSmoothImages = True
    
        #Based on the fluoro mask, i determined the following self.region of interest 
        # (row,column) format for numpy array:
        self.reg.SetRegionOfInterest(roi)
    
#===============================================================================
#    Start Metric evaluation loop
#==============================================================================
    def run_test(self,method='gc',tag='',starting_pose=[0,0,0,0,0,0]):
        
        self.reg.SetImageMetric(method)
        # Set the scale (reduces resolution of volume and fixed image)
        #fixedImage.ShrinkImage(level)
        #volume.ShrinkVolume(level)
        
        # Offset the poses by the optional starting pose
        all_poses = np.array([self.trans_range + starting_pose[0],
                              self.trans_range + starting_pose[1],
                              self.trans_range + starting_pose[2],
                              self.angle_range + starting_pose[3],
                              self.angle_range + starting_pose[4],
                              self.angle_range + starting_pose[5]]).transpose()
        
        num_evals = len(all_poses)
        val = np.zeros((num_evals,6))
        pylab.ion()
        fig = pylab.figure()
        fig.canvas.setFixedSize(800,900)
        fig.set_size_inches(8,10)
        fig.subplots_adjust(wspace=0.35,hspace=0.4)
        ax = [range(6)]
        
        for dim in range(6):
            # loop through each dimension of the pose (Tx, Ty, Tz, Rx, Ry, Rz)
            # only want to vary one parameter at a time, all others are set to 0 using a mask
            pose_mask = np.zeros((num_evals,6))
            pose_mask[:,dim]=1
            poses = all_poses*pose_mask
            for i,pose in enumerate( poses ):
                val[i,dim] = self.reg.GetValue(pose)
                print "    Metric: %s, Iteration: %i, Metric Value: %f" % \
                      (self.reg._metricParams['method'],i,val[i,dim])
                #print "    Pose: ", pose
            
            #plot results
            # Create a subplot for each dimension
            ax = fig.add_subplot('32'+str(dim+1))
            ax.plot(all_poses[:,dim],val[:,dim])
            ax.set_title(self.reg._metricParams['method']+' Dim:'+str(dim))
            ax.set_xlabel('Pose')
            ax.set_ylabel('Metric Value')
            ax.set_xlim(all_poses[0,dim],all_poses[-1,dim])
            pylab.draw()
        
        resultsFolder = 'C:/Users/bryan/bryan-code/trunk/Results/'
        save_name = resultsFolder+method.upper()+'_'+tag+'_'+'_metric_function.png'
        pylab.draw()
        pylab.savefig(save_name)
    


if __name__ == "__main__":
    #method_list = ['nmi', 'mi', 'ncc', 'ecc', 'mse', 'ncc', 'gd', 'gc']
    method_list = ['nc']
    tag = 'ITK_NCmetric'    # to label the resulting image that will be saved.
    for method in method_list:
        test = Test()
        # if you want to explore the region around a optimized solution
        #starting_pose=[0,-0.307422199545,1.12905408588,-0.0105103702415,0.00886278707541,0.000212203903439]
        test.run_test(method,tag)
        print "Done ... ", method
        
        